from .base_options import BaseOptions


class TrainOptions(BaseOptions):
    """
    处理训练模型时的命令行参数类，继承自BaseOptions
    :author: chaoshangcs(github), qms
    """

    def __init__(self):
        super().__init__()
        self.parser.add_argument('--is_train', type=bool, default=True,
                                 help='whether to train or test')
        self.parser.add_argument('--train_annotation_file', type=str, default='./data/msvd-qa/train_qa.json',
                                 help='annotation file path')
        self.parser.add_argument('--val_annotation_file', type=str, default='./data/msvd-qa/val_qa.json',
                                 help='val annotation file path')
        self.parser.add_argument('--mode', type=str, default='train',
                                 help='train or test or val')
        self.parser.add_argument('--train_question_pt', type=str, default='./data/msvd-qa/msvd-qa_train_questions.pt',
                                 help='train question pt file path')
        self.parser.add_argument('--val_question_pt', type=str, default='./data/msvd-qa/msvd-qa_val_questions.pt',
                                    help='val question pt file path')
        self.parser.add_argument('--lr', type=float, default=0.0001,
                                 help='initial learning rate for adam')
        self.parser.add_argument('--beta1', type=float, default=0.5, help='momentum term of adam')
        self.parser.add_argument('--max_epochs', type=int, default=50,
                                 help='number of epochs to train for')
        self.parser.add_argument('--stage_one', type=bool, default=False,
                                 help='whether to train stage one')
        self.parser.add_argument('--vision_dim', type=int, default=2048,
                                    help='vision dim')
        self.parser.add_argument('--module_dim', type=int, default=768,
                                    help='module dim')
        self.parser.add_argument('--word_dim', type=int, default=300,
                                    help='word dim')
        self.parser.add_argument('--graph_module', type=str, default='GAT',
                                    help='graph module')
        self.parser.add_argument('--graph_layers', type=int, default=1,
                                    help='graph layers')
        self.parser.add_argument('--unit_layers', type=int, default=2,
                                    help='unit layers')
        self.parser.add_argument('--restore', type=bool, default=False,
                                    help='whether to restore')
        self.parser.add_argument('--restore_path', type=str, default='./results/standard2',
                                    help='restore path')
        # 在这里定义其他命令行参数
